from langchain_core.prompts import ChatPromptTemplate
from langchain.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field
from patagent.llm import PatsnapLLM
from patagent.constant import *


llm = PatsnapLLM()

class CorporateAwardList(BaseModel):
    """The corporate list of the award."""
    award_name: str = Field(description="The name of the award shown in the image")
    corporates: list = Field(description="The list of the corporates shown in the image")

parser = PydanticOutputParser(pydantic_object=CorporateAwardList)

# Prompt
prompt = ChatPromptTemplate.from_messages([
    ("system", "Return the requested response object in Chinese. \n'{format_instructions}'\n"),
    ("human", [
        {
            "type": "image_url",
            "image_url": {"url": "{image_url}"},
        }
    ])
])

# Chain
ocr_chain = prompt | llm | parser